"""
-----------------------------------------------------------------------------------------
安装点刚度后处理程序 V0.4
作者：魏树森
版本释放日期：2024.07.08
-----------------------------------------------------------------------------------------
V0.4更新说明：
1.使用彩色字体，增加LOAD_SET ID显示，便于用户查找错误信息
2.如果F06文件找不到节点位移信息，提示用户

V0.3更新说明：
1.修复不同工况使用同一个load导致的程序崩溃，并提示用户可能头文件设定可能存在错误（多数情况是用户工况设定错误）

V0.2更新说明：
1.修复windows下程序执行的当前路径可能不是py文件存放路径，导致程序无法找到f06文件的问题
2.修复全局坐标系为null问题，将局部坐标系id定义为0
3.将方向输出由1、2、3改为全局坐标系方向输出为X 、Y 、Z ，局部坐标系输出为X'、Y'、Z'
4.完善程序提示信息

V0.1说明：
ost.py mains output stiffness table

"""
import glob
import csv
import math
import os
if os.name == 'nt':
    os.system("")

# Color类，用于控制台文本颜色、下划线、加粗等输出
class Color:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'  # 重置为默认颜色
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'


# SUBCASE类，用于整理关键信息
class Subcase:
    def __init__(self, sub_id, load_id, node_id, coord_id, force, direction, displacement, stiffness):
        self.sub_id = sub_id  # 用于获取SUBCASE编号
        self.load_id = load_id  # 用于获取该SUBCASE下的LOAD_SET编号
        self.node_id = node_id  # 用于获取该SUBCASE下的加载点的NODE ID
        self.coord_id = coord_id  # 用于获取该SUBCASE下的加载点坐标系ID，其中0为全局坐标系
        self.force = force  # 用于获取该SUBCASE下加载点的FORCE,force=FORCE*N
        self.direction = direction  # 用于获取该force的方向，其中1、2、3分别对应X、Y、Z
        self.displacement = displacement  # 用于获取该SUBCASE下的加载点位移
        self.stiffness = stiffness  # 用于计算该SUBCASE下的加载点的刚度

    # 类方法，用于在控制台打印表格，int(self.coord_id != 0)用于快速将非0整数转为1
    def show(self):
        print(Color.UNDERLINE + "|" + str(self.sub_id).center(8) + "|" + str(self.load_id).center(8) + "|" + str(self.node_id).center(8) + "|" + str(self.coord_id).center(8) + "|" + str(self.force).center(8) + "|" + ("XYZ"[self.direction - 1] + " '"[int(self.coord_id != 0)]).center(5) + "|" +str("{:.4f}".format(self.displacement)).center(10) + "|" +str("{:.2f}".format(self.stiffness)).center(10) + "|" + Color.ENDC)


# 进度条函数
def print_progress_bar(iteration, total, prefix='', suffix='', decimals=1, length=100, fill='█', print_end="\r"):
    percent = ("{0:." + str(decimals) + "f}").format(100 * (iteration / float(total)))
    filled_length = int(length * iteration // total)
    bar = fill * filled_length + '-' * (length - filled_length)
    print(Color.OKGREEN + f'\r{prefix} |{bar}| {percent}% {suffix}', end=print_end + Color.ENDC)
    # Print New Line on Complete
    if iteration == total:
        print()


# 该函数可以确保当前路径为py文件所在的路径
os.chdir(os.path.dirname(os.path.abspath(__file__)))

# 获取当前目录下的所有F06文件，并编号，如果F06文件只有一个，程序不询问用户，直接运行；如果超过一个，提醒用户输入编号后操作，避免手动输入造成错误
f06_files = glob.glob('*.f06')
for idx, file_path in enumerate(f06_files, start=1):
    print(Color.OKBLUE + f"f06文件编号：{idx}: {file_path}" + Color.ENDC)
if len(f06_files) == 0:
    print(Color.FAIL + "错误：当前文件夹内未发现F06文件" + Color.ENDC)
    input("按 ENTER 键退出")
elif len(f06_files) == 1:
    user_input = 1
else:
    user_input = int(
        input(Color.WARNING + "提示：发现多个F06文件，请输入要操作的文件编号后按回车键执行程序：" + Color.ENDC))
if 1 <= user_input <= len(f06_files):
    file_to_read = f06_files[user_input - 1]
    f = open(file_to_read, 'r')
else:
    print(Color.FAIL + "文件编号输入错误！" + Color.ENDC)
    input("按 ENTER 键退出")

# 输出提示信息，并初始化进度条
print(Color.OKBLUE + "正在从F06文件中获取信息:" + Color.ENDC)
print(Color.WARNING + "提示: 若进度条长时间未更新，请检查计算卡片是否设置了 DISPLACEMENT=ALL " + Color.ENDC)
print_progress_bar(0, 100, prefix='进度：', suffix='完成', length=50)

# 用于获取F06文件中所有的SUBCASE及其对应的LOAD_SET ID，创建SUBCASE的列表并初始化，获取完成后中断读取F06文件
subcases = []
subcaseID = 0
for line in f:
    if "SUBCASE" == line[25:32]:
        cases = line.split('SUBCASE')
        if len(cases) == 2:
            subcaseID = int(cases[1].strip())
    if "LOAD" in line:
        loads = line.split("LOAD")
        if len(loads) == 2 and "=" in loads[1]:
            loads = line.split('=')
            if len(loads) == 2:
                loadID = int(loads[1].strip())
                subcases.append(Subcase(subcaseID, loadID, 0, 0, 0.0, 0, 0.0, 0.0))
    if "S O R T E D   B U L K   D A T A   E C H O" in line:
        break

# 继续读取F06文件，获取对应的SUBCASE中的载荷、节点编号、坐标、方向信息，获取完成后中断读取
loads = []
for line in f:
    if 'FORCE' in line:
        loadID = int(line[38:46])
        nodeID = int(line[46:54])
        if line[54:62] == '        ':
            coord_id = 0
        else:
            coord_id = int(line[54:62])
        n1 = float(line[70:78])
        n2 = float(line[78:86])
        n3 = float(line[86:94])
        n = math.sqrt(n1 * n1 + n2 * n2 + n3 * n3)
        force = float(line[62:70]) * n
        direction = round(abs(n1 / n) + abs(n2 / n) * 2 + abs(n3 / n) * 3)
        counter = 0  # 用于计数
        loads.append([loadID, nodeID, coord_id, force, direction, counter])
    if "  M O D E L   S U M M A R Y   " in line:
        break

for sc in subcases:
    for ls in loads:
        if sc.load_id == ls[0]:
            sc.node_id = ls[1]
            sc.coord_id = ls[2]
            sc.force = ls[3]
            sc.direction = ls[4]
            ls[5] += 1
            break

# 继续读取F06文件，获取对应的NODE点的位移值
flag = 0  # 用于标记，如果flag=1表示已经读取完信息，程序将跳转到下一个subcase继续读取信息
i = 0  # 用于进度条
for sc in subcases:
    for line in f:
        if "      " + str(sc.node_id) + "      G     " in line and line[20] == 'G':
            sc.displacement = abs(float(line[11 + 15 * sc.direction:26 + 15 * sc.direction]))
            sc.stiffness = abs(sc.force / sc.displacement)
            i += 1
            print_progress_bar(i, len(subcases), prefix='进度：', suffix='完成', length=50)
            flag = 1
        if 'SUBCASE ' in line and 'SUBCASE ' + str(sc.sub_id) not in line and flag == 1:
            flag = 0
            break
f.close()

# 打印表头和表格
if i > 0:
    print(Color.OKBLUE + "\n完成：F06文件读取完毕，打印表格如下：" + Color.ENDC)
    print(
        Color.UNDERLINE + """                                                                          \n|SubCS ID|Load ID |Node ID |Coord ID|Force:N |Diret| Disp:mm  |Stiff:N/mm|""" + Color.ENDC)
    for sc in subcases:
        sc.show()

    # 在同目录下输出output.csv文件，csv文件为逗号分隔的文本文件，可以用EXCEL打开并保存到EXCEL表中编辑
    filename = 'output.csv'
    with open(filename, mode='w', newline='') as file:
        writer = csv.writer(file)
        header = ['SUBCASE ID', 'LOAD_SET ID', 'Node ID', 'Coord ID', 'Force(N)', 'Direction', 'Disp (mm)',
                  'Stiff (N/mm)']
        writer.writerow(header)
        for sc in subcases:
            writer.writerow(
                [sc.sub_id, sc.load_id, sc.node_id, sc.coord_id, sc.force,
                 ("XYZ"[sc.direction - 1] + " '"[int(sc.coord_id != 0)]).center(8),
                 sc.displacement, sc.stiffness])

    # 程序执行完毕，打印提示信息
    print()
    for ls in loads:
        if ls[5] > 1:
            print(
                Color.WARNING + f"警告： LOAD_SET {ls[0]} 被调用了 {ls[5]} 次！请检查Header文件是否正确，若正确，请忽略提示。 " + Color.ENDC)
    print(Color.OKGREEN + f"文件 {filename} 已创建并写入数据。" + Color.ENDC)
else:
    print(Color.FAIL + "错误：无法从F06文件中读取到节点位移信息！                                                      " + Color.ENDC)
    print(Color.WARNING + "提示：请确保计算模型中的位移输出设置为：" + Color.OKBLUE + " DISPLACEMENT=ALL " + Color.ENDC)
    print(
        Color.WARNING + "提示：如果计算模型中的位移输出设置为：" + Color.FAIL + "DISPLACEMENT(PLOT )=ALL" + Color.WARNING + " 将导致F06文件无位移结果写入！" + Color.ENDC)
if os.name == 'nt':
    input("程序执行完毕，按 ENTER 键退出!")
